package com.hexiang.ratelimit.store;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Slf4j
public class BucketStore {

    private final Map<String, Bucket> cache = new ConcurrentHashMap<>();

    public Bucket resolveBucket(String apiKey, long capacity, long tokens, long seconds) {
        return cache.computeIfAbsent(apiKey, key -> newBucket(capacity, tokens, seconds));
    }

    private Bucket newBucket(long capacity, long tokens, long seconds) {
        return Bucket4j.builder()
                .addLimit(Bandwidth.classic(capacity, Refill.intervally(tokens, Duration.ofSeconds(seconds))))
                .build();
    }

}